# 首先设置环境，在导入其他模块之前
import os
import sys
import logging

# 完全禁用所有代理设置
proxy_vars = ["ALL_PROXY", "HTTP_PROXY", "HTTPS_PROXY", "NO_PROXY", 
             "all_proxy", "http_proxy", "https_proxy", "no_proxy"]
for var in proxy_vars:
    if var in os.environ:
        del os.environ[var]

# 禁用litellm代理
os.environ["LITELLM_PROXY_ENABLED"] = "False"

# 为 httpx 库设置全局变量
os.environ["HTTPX_SKIP_PROXY_ENV_VARS"] = "1"

# 强制清除代理设置
os.environ["NO_PROXY"] = "*"

# 设置DeepSeek API密钥
os.environ["DEEPSEEK_API_KEY"] = "sk-99107bf855e349449ae5fa159f56187b"

# 现在开始导入其他模块
from google.protobuf.any import is_type
from langchain_openai import ChatOpenAI, AzureChatOpenAI
from langchain_deepseek import ChatDeepSeek
from langchain_google_genai import ChatGoogleGenerativeAI
from pydantic import SecretStr
from src.llms.litellm_v2 import ChatLiteLLMV2 as ChatLiteLLM
from src.config import load_yaml_config
from typing import Optional, Union
from litellm import LlmProviders
from pathlib import Path
from typing import Dict, Any

from litellm import CustomStreamWrapper, completion

from src.config import (
    REASONING_MODEL,
    REASONING_BASE_URL,
    REASONING_API_KEY,
    BASIC_MODEL,
    BASIC_BASE_URL,
    BASIC_API_KEY,
    VL_MODEL,
    VL_BASE_URL,
    VL_API_KEY,
    AZURE_API_BASE,
    AZURE_API_KEY,
    AZURE_API_VERSION,
    BASIC_AZURE_DEPLOYMENT,
    VL_AZURE_DEPLOYMENT,
    REASONING_AZURE_DEPLOYMENT,
)
from src.config.agents import LLMType



def create_openai_llm(
    model: str,
    base_url: Optional[str] = None,
    api_key: Optional[str] = None,
    temperature: float = 0.0,
    **kwargs,
) -> ChatOpenAI:
    """
    Create a ChatOpenAI instance with the specified configuration
    """
    # Only include base_url in the arguments if it's not None or empty
    llm_kwargs = {"model": model, "temperature": temperature, **kwargs}

    if base_url:  # This will handle None or empty string
        llm_kwargs["base_url"] = base_url

    if api_key:  # This will handle None or empty string
        llm_kwargs["api_key"] = api_key

    return ChatOpenAI(**llm_kwargs)


def create_deepseek_llm(
    model: str,
    base_url: Optional[str] = None,
    api_key: Optional[str] = None,
    temperature: float = 0.0,
    **kwargs,
) -> ChatDeepSeek:
    """
    Create a ChatDeepSeek instance with the specified configuration
    """
    # 确保重新加载.env文件
    from dotenv import load_dotenv
    load_dotenv(override=True)
    
    # 如果提供了api_key参数，优先使用参数值
    if api_key:
        os.environ["DEEPSEEK_API_KEY"] = api_key
    
    # 检查环境变量中是否有API密钥
    if "DEEPSEEK_API_KEY" not in os.environ or not os.environ["DEEPSEEK_API_KEY"]:
        # 打印当前所有环境变量，便于调试
        print("\n\n当前环境变量列表:")
        for key, value in os.environ.items():
            if "key" in key.lower() or "api" in key.lower():
                print(f"{key}={value[:5]}{'*' * 10}" if value else f"{key}=None")
        raise ValueError("DEEPSEEK_API_KEY未设置。请提供有效的API密钥或设置环境变量")
    
    # 准备参数
    llm_kwargs = {"model": "deepseek-chat", "temperature": temperature, "max_tokens": None, "timeout": None}
    
    # 如果有base_url，添加到参数中
    if base_url:
        llm_kwargs["api_base"] = base_url
        
    # 创建实例
    try:
        return ChatDeepSeek(**llm_kwargs)
    except Exception as e:
        raise ValueError(f"DeepSeek初始化失败: {e}")


def create_azure_llm(
    azure_deployment: str,
    azure_endpoint: str,
    api_version: str,
    api_key: str,
    temperature: float = 0.0,
) -> AzureChatOpenAI:
    """
    create azure llm instance with specified configuration
    """
    return AzureChatOpenAI(
        azure_deployment=azure_deployment,
        azure_endpoint=azure_endpoint,
        api_version=api_version,
        api_key=api_key,
        temperature=temperature,
    )


def create_gemini_llm(
    model: str,
    api_key: Optional[str] = None,
    temperature: float = 0.0,
    **kwargs,
) -> ChatGoogleGenerativeAI:
    """
    创建 Gemini 模型实例，特别适用于视觉处理任务
    """
    # 检查环境变量中是否有API密钥，没有则使用传入的api_key
    gemini_api_key = api_key or os.environ.get("GEMINI_API_KEY")
    
    if not gemini_api_key:
        raise ValueError("Gemini API密钥未设置。请提供有效的API密钥或设置环境变量")
    
    return ChatGoogleGenerativeAI(
        model=model,
        temperature=temperature,
        api_key=SecretStr(gemini_api_key),
        convert_system_message_to_human=True,
        **kwargs
    )

def create_litellm_model(
    model: str,
    base_url: Optional[str] = None,
    api_key: Optional[str] = None,
    temperature: float = 0.0,
    **kwargs,
) -> ChatLiteLLM:
    """
    Support various different model's through LiteLLM's capabilities.
    """

    llm_kwargs = {"model": model, "temperature": temperature, **kwargs}

    if base_url:  # This will handle None or empty string
        llm_kwargs["api_base"] = base_url

    if api_key:  # This will handle None or empty string
        llm_kwargs["api_key"] = api_key

    return ChatLiteLLM(**llm_kwargs)


# Cache for LLM instances
_llm_cache: dict[LLMType, Union[ChatOpenAI, ChatDeepSeek, AzureChatOpenAI, ChatLiteLLM, ChatGoogleGenerativeAI]] = (
    {}
)


def is_litellm_model(model_name: str) -> bool:
    """
    Check if the model name indicates it should be handled by LiteLLM.

    Args:
        model_name: The name of the model to check

    Returns:
        bool: True if the model should be handled by LiteLLM, False otherwise
    """
    return (
        model_name
        and "/" in model_name
        and model_name.split("/")[0] in [p.value for p in LlmProviders]
    )


def _create_llm_use_env(
    llm_type: LLMType,
) -> Union[ChatDeepSeek, ChatGoogleGenerativeAI, ChatOpenAI]:
    # 显式设置 API 密钥环境变量
    if llm_type == "reasoning":
        api_key = REASONING_API_KEY
        model = REASONING_MODEL
        base_url = REASONING_BASE_URL
    elif llm_type == "basic":
        api_key = BASIC_API_KEY
        model = BASIC_MODEL
        base_url = BASIC_BASE_URL
    elif llm_type == "vision":
        api_key = VL_API_KEY
        model = VL_MODEL
        base_url = VL_BASE_URL
    else:
        raise ValueError(f"Unknown LLM type: {llm_type}")
    
    # 判断模型类型
    if model and "gemini" in model.lower():
        # 使用 Gemini 模型
        logging.info(f"使用Gemini模型: {model}")
        return create_gemini_llm(
            model=model,
            api_key=api_key,
            temperature=0.0,
        )
    elif model and ("gpt" in model.lower() or "openai" in model.lower()):
        # 使用 OpenAI 模型
        logging.info(f"使用OpenAI模型: {model}")
        return create_openai_llm(
            model=model,
            base_url=base_url,
            api_key=api_key,
            temperature=0.0,
        )
    else:
        # 使用 DeepSeek 模型（默认）
        logging.info(f"使用DeepSeek模型: {model}")
        os.environ["DEEPSEEK_API_KEY"] = api_key if api_key else ""
        return create_deepseek_llm(
            model=model,
            base_url=base_url,
            api_key=api_key,
        )


def _create_llm_use_conf(llm_type: LLMType, conf: Dict[str, Any]) -> ChatLiteLLM:
    llm_type_map = {
        "reasoning": conf.get("REASONING_MODEL"),
        "basic": conf.get("BASIC_MODEL"),
        "vision": conf.get("VISION_MODEL"),
    }
    llm_conf = llm_type_map.get(llm_type)
    if not llm_conf:
        raise ValueError(f"Unknown LLM type: {llm_type}")
    if not isinstance(llm_conf, dict):
        raise ValueError(f"Invalid LLM Conf: {llm_type}")
    return ChatLiteLLM(**llm_conf)


def get_llm_by_type(
    llm_type: LLMType,
) -> Union[ChatOpenAI, ChatDeepSeek, AzureChatOpenAI, ChatLiteLLM, ChatGoogleGenerativeAI]:
    """
    Get LLM instance by type. Returns cached instance if available.
    """
    if llm_type in _llm_cache:
        return _llm_cache[llm_type]

    conf = load_yaml_config(
        str((Path(__file__).parent.parent.parent / "conf.yaml").resolve())
    )
    use_conf = conf.get("USE_CONF", False)
    if use_conf:
        llm = _create_llm_use_conf(llm_type, conf)
    else:
        llm = _create_llm_use_env(llm_type)

    _llm_cache[llm_type] = llm
    return llm


# Initialize LLMs for different purposes - now these will be cached
reasoning_llm = get_llm_by_type("reasoning")
basic_llm = get_llm_by_type("basic")
vl_llm = get_llm_by_type("vision")


if __name__ == "__main__":
    # stream = reasoning_llm.stream("what is mcp?")
    # full_response = ""
    # for chunk in stream:
    #     full_response += chunk.content
    # print(full_response)

    print(basic_llm.invoke("Hello"))
    # print(vl_llm.invoke("Hello"))
